import java.util.ArrayList;
import java.util.LinkedList;
import java.util.List;
import java.util.Queue;

/**
 * 637. 二叉树的层平均值
 * https://leetcode-cn.com/problems/average-of-levels-in-binary-tree/
 */
public class Solutions_637 {
    public static void main(String[] args) {
        TreeNode root = new TreeNode(3);
        root.left = new TreeNode(9);
        root.right = new TreeNode(20);
        root.right.left = new TreeNode(15);
        root.right.right = new TreeNode(7);  // output: {3, 14.5, 11}

//        List<Double> result = averageOfLevels(root);
        List<Double> result = averageOfLevels2(root);
        System.out.println(result);
    }

    // 深度优先搜索-dfs
    public static List<Double> averageOfLevels(TreeNode root) {
        // 1.记录每一层的节点值总和
        List<Double> sum = new ArrayList<>();
        // 2.记录每一层的节点数量
        List<Integer> nodeCount = new ArrayList<>();
        // 3.遍历二叉树
        averageOfLevels_dfs(root, 0, sum, nodeCount);

        for (int i = 0; i < sum.size(); i++) {
            // 4.计算平均数
            double averageNum = sum.get(i) / nodeCount.get(i);
            sum.set(i, averageNum);
        }
        return sum;
    }

    // 先序遍历（根左右）二叉树
    public static void averageOfLevels_dfs(TreeNode root, int level, List<Double> sum,
                                           List<Integer> nodeCount) {
        if (root == null) {
            return;
        }
        if (level >= sum.size()) {
            // 遍历到第 level 层的第一个节点时，必须往 list 中添加默认值，否则 get(level) 时报错
            sum.add(0.0);
            nodeCount.add(0);
        }
        // level 层加上遍历的节点值，level 层上节点个数加 1
        sum.set(level, sum.get(level) + root.val);
        nodeCount.set(level, nodeCount.get(level) + 1);

        // 遍历左右子树
        averageOfLevels_dfs(root.left, level + 1, sum, nodeCount);
        averageOfLevels_dfs(root.right, level + 1, sum, nodeCount);
    }

    // 广度优先搜索-bfs
    public static List<Double> averageOfLevels2(TreeNode root) {
        List<Double> res = new ArrayList<>();
        if (root == null) {
            return res;
        }
        // 1.使用队列记录遍历的节点对象
        Queue<TreeNode> queue = new LinkedList<>();
        // 2.添加根节点
        queue.add(root);
        // 3.循环队列，完成遍历二叉树
        while (!queue.isEmpty()) {
            int size = queue.size();
            double sum = 0.0;
            for (int i = 0; i < size; i++) {
                TreeNode node = queue.poll();
                sum += node.val;
                if (node.left != null) {
                    queue.add(node.left);
                }
                if (node.right != null) {
                    queue.add(node.right);
                }
            }
            // 4.计算每一层的平均值
            res.add(sum / size);
        }
        return res;
    }
}
